package com.hyw.cm.protocal.slow;

import java.util.Objects;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import io.netty.util.HashedWheelTimer;
import io.netty.util.Timeout;
import io.netty.util.Timer;
import io.netty.util.TimerTask;

/**
 * 针对慢读攻击做一个最大可容忍的超时判断功能。如果超时时间内，包无法收齐，则强制中断当前连接，以提升连接资源回收
 * 
 * @author Hongyu
 */
public class SlowAttackLengthFieldBasedFrameDecoder extends LengthFieldBasedFrameDecoder {

  private static final Logger log = LoggerFactory.getLogger(SlowAttackLengthFieldBasedFrameDecoder.class);

  /**
   * 慢读攻击容忍时间为5秒
   */
  private static final int MIN_SLOW_TIMEOUT = 5;
  private static final int MAX_SLOW_TIMEOUT = 60;

  private static final Timer timer = new HashedWheelTimer();

  private volatile Timeout current;

  private final int slowTimeout;

  public SlowAttackLengthFieldBasedFrameDecoder(int maxFrameLength, int lengthFieldOffset, int lengthFieldLength,
      int lengthAdjustment, int initialBytesToStrip, boolean failFast) {
    this(maxFrameLength, lengthFieldOffset, lengthFieldLength, lengthAdjustment, initialBytesToStrip, failFast, 0);
  }

  /**
   * @param maxFrameLength
   * @param lengthFieldOffset
   * @param lengthFieldLength
   * @param lengthAdjustment
   * @param initialBytesToStrip
   * @param failFast
   * @param slowAttackTimeoutInSeconds
   *          参数值小于或者等于0，则不启动慢读攻击检测。如果参数值小于5，则使用默认的超时时间。如果大于5，则按入参值计算。该值的最大可接受值，不得超过60秒，否则，会抛出{@link IllegalArgumentException}
   */
  public SlowAttackLengthFieldBasedFrameDecoder(int maxFrameLength, int lengthFieldOffset, int lengthFieldLength,
      int lengthAdjustment, int initialBytesToStrip, boolean failFast, int slowAttackTimeoutInSeconds) {
    super(maxFrameLength, lengthFieldOffset, lengthFieldLength, lengthAdjustment, initialBytesToStrip, failFast);
    Assert.isTrue(slowAttackTimeoutInSeconds <= MAX_SLOW_TIMEOUT, "取值范围[1,60]");
    this.slowTimeout = Math.max(MIN_SLOW_TIMEOUT, slowAttackTimeoutInSeconds);
  }

  @Override
  protected Object decode(ChannelHandlerContext ctx, ByteBuf in) throws Exception {
    Object data = super.decode(ctx, in);

    if (Objects.isNull(data) && this.slowTimeout > 0 && Objects.isNull(current)) {
      current = timer.newTimeout(new TimerTask() {
        @Override
        public void run(Timeout timeout) throws Exception {
          log.warn("检测到slow attack. ctx:{}", ctx.channel().id());
          if (!timeout.isCancelled() && ctx.channel().isOpen()) {
            ctx.channel().close();
          }
        }
      }, slowTimeout, TimeUnit.SECONDS);
    }

    if (Objects.nonNull(data)) {
      this.cancelTimeout();
    }

    return data;
  }

  @Override
  public void channelInactive(ChannelHandlerContext ctx) throws Exception {
    this.cancelTimeout();
    super.channelInactive(ctx);
  }

  private void cancelTimeout() {
    if (Objects.nonNull(current)) {
      current.cancel();
      current = null;
    }
  }

}